// STL
#include <fstream>                         // std::ofstream
#include <string>                          // std::string
#include <vector>                          // std::vector<>

// Boost
#include <boost/archive/text_oarchive.hpp> // boost::archive::text_oarchive

// jScience
#include "jScience/linalg.hpp"             // Matrix<>

// stats++
#include "statsxx/dataset.hpp"             // Matrix_to_DataSet(), shuffle_data_set(), partition_data_set()
#include "statsxx/machine_learning/NeuralNet.hpp"      // NEURAL_NET
#include "statsxx/machine_learning/neural_network.hpp" // NeuralNet_EA_train()
#include "statsxx/optimization/EA.hpp"                 // EAParam



//
// DESC: Trains a MLP.
//
NEURAL_NET train_MLP(
                     NEURAL_NET            mlp,
                     // -----
                     const int             method,
                     //------
                     const int             nepoch_min,
                     const int             nepoch_max,
                     // -----
                     const bool            early_stopping,
                     // -----
                     const double          lr,
                     const double          lr_min,
                     const double          lr_max,
                     // -----
                     const double          momentum,
                     // -----
                     const double          weight_penalty,
                     // -----
                     const double          qrprop_u,
                     const double          qrprop_d,
                     // -----
                     const double          scg_lambda,
                     const double          scg_sigma,
                     const double          scg_convg_iterfrac,
                     const double          scg_rk_tol,
                     // -----
                     EAParam               EA_param,
                     // -----
                     const Matrix<double> &X_in,
                     const Matrix<double> &X_out,
                     // -----
                     const int             npts_per_batch,
                     // -----
                     const std::string     mlp_file
                     )
{
    const bool silent  = false;

    //=========================================================

    //=========================================================
    // INITIALIZATION
    //=========================================================

    DataSet dataset = Matrix_to_DataSet(
                                        X_in,
                                        X_out
                                        );

    shuffle_data_set(dataset);

    std::vector<DataSet> datasets = partition_data_set(
                                                       dataset
                                                       );

    //=========================================================
    // TRAIN
    //=========================================================

    if( method != 99 )
    {
        mlp.train(
                  method,
                  //------
                  nepoch_min,
                  nepoch_max,
                  // -----
                  early_stopping,
                  // -----
                  lr,
                  lr_min,
                  lr_max,
                  // -----
                  momentum,
                  // -----
                  weight_penalty,
                  // -----
                  qrprop_u,
                  qrprop_d,
                  // -----
                  scg_lambda,
                  scg_sigma,
                  scg_convg_iterfrac,
                  scg_rk_tol,
                  // -----
                  datasets[0],
                  datasets[1],
                  datasets[2],
                  // -----
                  npts_per_batch,
                  // -----
                  silent
                  );
    }
    else
    {
        EA_param.genome_size = mlp.get_weights().size();

        mlp = NeuralNet_EA_train(
                                 mlp,
                                 // -----
                                 EA_param,
                                 // -----
                                 datasets[0],
                                 datasets[1]
                                 );
    }

    //=========================================================
    // OUTPUT
    //=========================================================

    // SAVE MLP TO ARCHIVE
    {
        std::ofstream ofs(mlp_file, std::ios::out);

        boost::archive::text_oarchive oa(ofs);

        oa << mlp;
    }

    return mlp;
}
